#include <srs_app_statistic.hpp>

#include <unistd.h>
#include <sstream>
using namespace std;

#include <srs_rtmp_stack.hpp>
#include <srs_protocol_json.hpp>
#include <srs_protocol_kbps.hpp>
#include <srs_app_conn.hpp>
#include <srs_app_config.hpp>
#include <srs_kernel_utility.hpp>

int64_t srs_gvid = getpid();

int64_t srs_generate_id() {
    return srs_gvid++;
}

SrsStatisticVhost::SrsStatisticVhost() {
    id = srs_generate_id();

    kbps = new SrsKbps();
    kbps->set_io(nullptr, nullptr);

    nb_clients = 0;
    nb_streams = 0;
}

SrsStatisticVhost::~SrsStatisticVhost() {
    srs_freep(kbps);
}

int SrsStatisticVhost::dumps(stringstream& ss) {
    int ret = ERROR_SUCCESS;

    // dumps the config of vhost.
    bool hls_enabled = _srs_config->get_hls_enabled(vhost);
    bool enabled = _srs_config->get_vhost_enabled(vhost);

    ss << SRS_JOBJECT_START
       << SRS_JFIELD_ORG("id", id) << SRS_JFIELD_CONT
       << SRS_JFIELD_STR("name", vhost) << SRS_JFIELD_CONT
       << SRS_JFIELD_BOOL("enabled", enabled) << SRS_JFIELD_CONT
       << SRS_JFIELD_ORG("clients", nb_clients) << SRS_JFIELD_CONT
       << SRS_JFIELD_ORG("streams", nb_streams) << SRS_JFIELD_CONT
       << SRS_JFIELD_ORG("send_bytes", kbps->get_send_bytes()) << SRS_JFIELD_CONT
       << SRS_JFIELD_ORG("recv_bytes", kbps->get_recv_bytes()) << SRS_JFIELD_CONT
       << SRS_JFIELD_OBJ("kbps")
       << SRS_JFIELD_ORG("recv_30s", kbps->get_recv_kbps_30s()) << SRS_JFIELD_CONT
       << SRS_JFIELD_ORG("send_30s", kbps->get_send_kbps_30s())
       << SRS_JOBJECT_END << SRS_JFIELD_CONT
       << SRS_JFIELD_NAME("hls") << SRS_JOBJECT_START
       << SRS_JFIELD_BOOL("enabled", hls_enabled);

    if (hls_enabled) {
        ss << SRS_JFIELD_CONT;
        ss << SRS_JFIELD_ORG("fragment", _srs_config->get_hls_fragment(_srs_config->get_hls(vhost)));
    }
    ss << SRS_JOBJECT_END
       << SRS_JOBJECT_END;

    return ret;
}

SrsStatisticStream::SrsStatisticStream() {
    id = srs_generate_id();
    vhost = nullptr;
    active = false;
    connection_cid = -1;

    has_video = false;
    vcodec = SrsCodecVideoReserved;
    avc_profile = SrsAvcProfileReserved;
    avc_level = SrsAvcLevelReserved;

    has_audio = false;
    acodec = SrsCodecAudioReserved1;
    asample_rate = SrsCodecAudioSampleRateReserved;
    asound_type = SrsCodecAudioSoundTypeReserved;
    aac_object = SrsAacObjectTypeReserved;

    kbps = new SrsKbps();
    kbps->set_io(nullptr, nullptr);

    nb_clients = 0;
    nb_frames = 0;
}

SrsStatisticStream::~SrsStatisticStream() {
    srs_freep(kbps);
}

int SrsStatisticStream::dumps(stringstream& ss) {
    int ret = ERROR_SUCCESS;

    ss << SRS_JOBJECT_START
       << SRS_JFIELD_ORG("id", id) << SRS_JFIELD_CONT
       << SRS_JFIELD_STR("name", stream) << SRS_JFIELD_CONT
       << SRS_JFIELD_STR("stream", stream) << SRS_JFIELD_CONT
       << SRS_JFIELD_STR("stream_id", stream_id) << SRS_JFIELD_CONT
       << SRS_JFIELD_ORG("vhost", vhost->id) << SRS_JFIELD_CONT
       << SRS_JFIELD_STR("vhost_name", vhost->vhost) << SRS_JFIELD_CONT
       << SRS_JFIELD_STR("app", app) << SRS_JFIELD_CONT
       << SRS_JFIELD_ORG("live_ms", srs_get_system_time_ms()) << SRS_JFIELD_CONT
       << SRS_JFIELD_ORG("clients", nb_clients) << SRS_JFIELD_CONT
       << SRS_JFIELD_ORG("frames", nb_frames) << SRS_JFIELD_CONT
       << SRS_JFIELD_ORG("send_bytes", kbps->get_send_bytes()) << SRS_JFIELD_CONT
       << SRS_JFIELD_ORG("recv_bytes", kbps->get_recv_bytes()) << SRS_JFIELD_CONT
       << SRS_JFIELD_OBJ("kbps")
       << SRS_JFIELD_ORG("recv_30s", kbps->get_recv_kbps_30s()) << SRS_JFIELD_CONT
       << SRS_JFIELD_ORG("send_30s", kbps->get_send_kbps_30s())
       << SRS_JOBJECT_END << SRS_JFIELD_CONT
       << SRS_JFIELD_OBJ("publish")
       << SRS_JFIELD_BOOL("active", active) << SRS_JFIELD_CONT
       << SRS_JFIELD_ORG("cid", connection_cid)
       << SRS_JOBJECT_END << SRS_JFIELD_CONT;

    if (!has_video) {
        ss  << SRS_JFIELD_NULL("video") << SRS_JFIELD_CONT;
    } else {
        ss  << SRS_JFIELD_NAME("video") << SRS_JOBJECT_START
            << SRS_JFIELD_STR("codec", srs_codec_video2str(vcodec)) << SRS_JFIELD_CONT
            << SRS_JFIELD_STR("profile", srs_codec_avc_profile2str(avc_profile)) << SRS_JFIELD_CONT
            << SRS_JFIELD_STR("level", srs_codec_avc_level2str(avc_level))
            << SRS_JOBJECT_END
            << SRS_JFIELD_CONT;
    }

    if (!has_audio) {
        ss  << SRS_JFIELD_NULL("audio");
    } else {
        ss  << SRS_JFIELD_NAME("audio") << SRS_JOBJECT_START
            << SRS_JFIELD_STR("codec", srs_codec_audio2str(acodec)) << SRS_JFIELD_CONT
            << SRS_JFIELD_ORG("sample_rate", (int)flv_sample_rates[asample_rate]) << SRS_JFIELD_CONT
            << SRS_JFIELD_ORG("channel", (int)asound_type + 1) << SRS_JFIELD_CONT
            << SRS_JFIELD_STR("profile", srs_codec_aac_object2str(aac_object))
            << SRS_JOBJECT_END;
    }

    ss << SRS_JOBJECT_END;

    return ret;
}

void SrsStatisticStream::publish(const std::string& streamid, int cid) {
    stream_id = streamid;
    connection_cid = cid;
    active = true;

    vhost->nb_streams++;
}

void SrsStatisticStream::close() {
    has_video = false;
    has_audio = false;
    active = false;

    vhost->nb_streams--;
}

SrsStatisticClient::SrsStatisticClient() {
    id = 0;
    streamid = 0;
    vhostid = 0;
    conn = nullptr;
    req = nullptr;
    type = SrsRtmpConnUnknown;
    create = srs_get_system_time_ms();
}

SrsStatisticClient::~SrsStatisticClient() {
}

int SrsStatisticClient::dumps(stringstream& ss) {
    int ret = ERROR_SUCCESS;

    ss << SRS_JOBJECT_START
       << SRS_JFIELD_ORG("id", id) << SRS_JFIELD_CONT
       << SRS_JFIELD_ORG("vhost", vhostid) << SRS_JFIELD_CONT
       << SRS_JFIELD_ORG("stream", streamid) << SRS_JFIELD_CONT
       << SRS_JFIELD_STR("ip", req->ip) << SRS_JFIELD_CONT
       << SRS_JFIELD_STR("pageUrl", req->pageUrl) << SRS_JFIELD_CONT
       << SRS_JFIELD_STR("swfUrl", req->swfUrl) << SRS_JFIELD_CONT
       << SRS_JFIELD_STR("tcUrl", req->tcUrl) << SRS_JFIELD_CONT
       << SRS_JFIELD_STR("url", req->get_stream_url()) << SRS_JFIELD_CONT
       << SRS_JFIELD_STR("type", srs_client_type_string(type)) << SRS_JFIELD_CONT
       << SRS_JFIELD_BOOL("publish", srs_client_type_is_publish(type)) << SRS_JFIELD_CONT
       << SRS_JFIELD_ORG("alive", srs_get_system_time_ms() - create)
       << SRS_JOBJECT_END;

    return ret;
}

SrsStatistic* SrsStatistic::_instance = new SrsStatistic();

SrsStatistic::SrsStatistic() {
    _generate_id = srs_generate_id();

    kbps = new SrsKbps();
    kbps->set_io(nullptr, nullptr);
}

SrsStatistic::~SrsStatistic() {
    srs_freep(kbps);

    if (true) {
        for (auto it = vhosts.begin(); it != vhosts.end(); it++) {
            auto vhost = it->second;
            srs_freep(vhost);
        }
    }
    if (true) {
        for (auto it = streams.begin(); it != streams.end(); it++) {
            auto stream = it->second;
            srs_freep(stream);
        }
    }
    if (true) {
        for (auto it = clients.begin(); it != clients.end(); it++) {
            auto client = it->second;
            srs_freep(client);
        }
    }

    vhosts.clear();
    rvhosts.clear();
    streams.clear();
    rstreams.clear();
}

SrsStatistic* SrsStatistic::instance() {
    return _instance;
}

SrsStatisticVhost* SrsStatistic::find_vhost(int64_t vid) {
    auto it = vhosts.find(vid);
    if (it != vhosts.end()) {
        return it->second;
    }
    return nullptr;
}

SrsStatisticStream* SrsStatistic::find_stream(int64_t sid) {
    auto it = streams.find(sid);
    if (it != streams.end()) {
        return it->second;
    }
    return nullptr;
}

vector<SrsStatisticStream*> SrsStatistic::find_stream(const std::string& app, const std::string& stream) {
    vector<SrsStatisticStream*> arr;
    for (auto it : streams) {
        auto pstr = it.second;
        if (pstr == nullptr) {
            continue;
        }
        if (pstr->app == app && pstr->stream == stream) {
            arr.push_back(pstr);
        }
    }
    return arr;
}

SrsStatisticClient* SrsStatistic::find_client(int cid) {
    auto it = clients.find(cid);
    if (it != clients.end()) {
        return it->second;
    }
    return nullptr;
}

vector<SrsStatisticClient*> SrsStatistic::find_client_by_streamid(int64_t sid) {
    vector<SrsStatisticClient*> arr;
    for (auto it : clients) {
        auto client = it.second;
        if (client == nullptr) {
            continue;
        }
        if (client->streamid == sid) {
            arr.push_back(client);
        }
    }
    return arr;
}

int SrsStatistic::on_video_info(SrsRequest* req, SrsCodecVideo vcodec, SrsAvcProfile avc_profile, SrsAvcLevel avc_level) {
    int ret = ERROR_SUCCESS;
    auto vhost = create_vhost(req);
    if (vhost == nullptr) {
        return ret;
    }
    auto stream = create_stream(vhost, req);
    if (stream == nullptr) {
        return ret;
    }
    stream->has_video = true;
    stream->vcodec = vcodec;
    stream->avc_profile = avc_profile;
    stream->avc_level = avc_level;
    return ret;
}

int SrsStatistic::on_audio_info(SrsRequest* req, SrsCodecAudio acodec, SrsCodecAudioSampleRate asample_rate,
                                SrsCodecAudioSoundType asound_type, SrsAacObjectType aac_object) {
    int ret = ERROR_SUCCESS;
    auto vhost = create_vhost(req);
    if (vhost == nullptr) {
        return ret;
    }
    auto stream = create_stream(vhost, req);
    if (stream == nullptr) {
        return ret;
    }
    stream->has_audio = true;
    stream->acodec = acodec;
    stream->asample_rate = asample_rate;
    stream->asound_type = asound_type;
    stream->aac_object = aac_object;
    return ret;
}

int SrsStatistic::on_video_frames(SrsRequest* req, int nb_frames) {
    int ret = ERROR_SUCCESS;
    auto vhost = create_vhost(req);
    if (vhost == nullptr) {
        return ret;
    }
    auto stream = create_stream(vhost, req);
    if (stream == nullptr) {
        return ret;
    }
    stream->nb_frames += nb_frames;
    return ret;
}

void SrsStatistic::on_stream_publish(SrsRequest* req, int cid) {
    auto vhost = create_vhost(req);
    if (vhost == nullptr) {
        return;
    }
    auto stream = create_stream(vhost, req);
    if (stream == nullptr) {
        return;
    }
    stream->publish(req->stream_id, cid);
}

void SrsStatistic::on_stream_close(SrsRequest* req) {
    auto vhost = create_vhost(req);
    if (vhost == nullptr) {
        return;
    }
    auto stream = create_stream(vhost, req);
    if (stream == nullptr) {
        return;
    }
    stream->close();
}

int SrsStatistic::on_client(int id, SrsRequest* req, SrsConnection* conn, SrsRtmpConnType type) {
    int ret = ERROR_SUCCESS;
    auto vhost = create_vhost(req);
    if (vhost == nullptr) {
        return ret;
    }
    auto stream = create_stream(vhost, req);
    if (stream == nullptr) {
        return ret;
    }
    // create client if not exists
    SrsStatisticClient* client = nullptr;
    if (clients.find(id) == clients.end()) {
        client = new SrsStatisticClient();
        client->id = id;
        client->streamid = stream->id;
        if (stream->vhost != nullptr) {
            client->vhostid = stream->vhost->id;
        }
        clients[id] = client;
    } else {
        client = clients[id];
    }

    // got client.
    client->conn = conn;
    client->req = req;
    client->type = type;
    stream->nb_clients++;
    vhost->nb_clients++;

    return ret;
}

void SrsStatistic::on_disconnect(int id) {
    auto it = clients.find(id);
    if (it == clients.end()) {
        return;
    }

    int64_t streamid = 0, vhostid = 0;
    auto client = it->second;
    if (client != nullptr) {
        streamid = client->streamid;
        vhostid = client->vhostid;
        srs_freep(client);
    }
    clients.erase(it);

    if (streamid > 0) {
        auto stat = SrsStatistic::instance();
        if (stat != nullptr) {
            auto stream = stat->find_stream(streamid);
            if (stream != nullptr) {
                stream->nb_clients--;

                auto vhost = stream->vhost;
                if (vhost == nullptr) {
                    vhost = stat->find_vhost(vhostid);
                }
                if (vhost != nullptr) {
                    vhost->nb_clients--;
                }
            }
        }
    }
}

void SrsStatistic::kbps_add_delta(SrsConnection* conn) {
    int id = conn->srs_id();
    if (clients.find(id) == clients.end()) {
        return;
    }

    auto client = clients[id];
    if (client == nullptr) {
        return;
    }

    // resample the kbps to collect the delta.
    conn->resample();

    // add delta of connection to kbps.
    // for next sample() of server kbps can get the stat.
    kbps->add_delta(conn);

    if (client->streamid > 0) {
        auto stat = SrsStatistic::instance();
        if (stat != nullptr) {
            auto stream = stat->find_stream(client->streamid);
            if (stream != nullptr) {
                if (stream->kbps != nullptr) {
                    stream->kbps->add_delta(conn);
                }
                auto vhost = stream->vhost;
                if (vhost == nullptr) {
                    vhost = stat->find_vhost(client->vhostid);
                }
                if (vhost != nullptr) {
                    vhost->kbps->add_delta(conn);
                }
            }
        }
    }

    // cleanup the delta.
    conn->cleanup();
}

SrsKbps* SrsStatistic::kbps_sample() {
    kbps->sample();
    if (true) {
        for (auto it = vhosts.begin(); it != vhosts.end(); it++) {
            auto vhost = it->second;
            if (vhost != nullptr && vhost->kbps != nullptr) {
                vhost->kbps->sample();
            }
        }
    }
    if (true) {
        for (auto it = streams.begin(); it != streams.end(); it++) {
            auto stream = it->second;
            if (stream != nullptr && stream->kbps != nullptr) {
                stream->kbps->sample();
            }
        }
    }
    return kbps;
}

int64_t SrsStatistic::generate_id() {
    return _generate_id;
}

int SrsStatistic::dumps_vhosts(stringstream& ss) {
    int ret = ERROR_SUCCESS;
    ss << SRS_JARRAY_START;

    for (auto it = vhosts.begin(); it != vhosts.end(); it++) {
        auto vhost = it->second;
        if (vhost == nullptr) {
            continue;
        }
        if (it != vhosts.begin()) {
            ss << SRS_JFIELD_CONT;
        }
        if ((ret = vhost->dumps(ss)) != ERROR_SUCCESS) {
            return ret;
        }
    }

    ss << SRS_JARRAY_END;
    return ret;
}

int SrsStatistic::dumps_streams(stringstream& ss) {
    int ret = ERROR_SUCCESS;
    ss << SRS_JARRAY_START;

    for (auto it = streams.begin(); it != streams.end(); it++) {
        auto stream = it->second;
        if (stream == nullptr) {
            continue;
        }
        if (it != streams.begin()) {
            ss << SRS_JFIELD_CONT;
        }
        if ((ret = stream->dumps(ss)) != ERROR_SUCCESS) {
            return ret;
        }
    }

    ss << SRS_JARRAY_END;
    return ret;
}

int SrsStatistic::dumps_clients(stringstream& ss, int start, int count) {
    int ret = ERROR_SUCCESS;
    ss << SRS_JARRAY_START;

    auto it = clients.begin();
    for (int i = 0; i < start + count && it != clients.end(); it++, i++) {
        if (i < start) {
            continue;
        }
        auto client = it->second;
        if (client == nullptr) {
            continue;
        }
        if (i != start) {
            ss << SRS_JFIELD_CONT;
        }
        if ((ret = client->dumps(ss)) != ERROR_SUCCESS) {
            return ret;
        }
    }

    ss << SRS_JARRAY_END;
    return ret;
}

SrsStatisticVhost* SrsStatistic::create_vhost(SrsRequest* req) {
    SrsStatisticVhost* vhost = nullptr;

    // create vhost if not exists.
    if (rvhosts.find(req->vhost) == rvhosts.end()) {
        vhost = new SrsStatisticVhost();
        vhost->vhost = req->vhost;
        rvhosts[req->vhost] = vhost;
        vhosts[vhost->id] = vhost;
        return vhost;
    }

    vhost = rvhosts[req->vhost];
    return vhost;
}

SrsStatisticStream* SrsStatistic::create_stream(SrsStatisticVhost* vhost, SrsRequest* req) {
    std::string url = req->get_stream_url();
    SrsStatisticStream* stream = nullptr;

    // create stream if not exists.
    if (rstreams.find(url) == rstreams.end()) {
        stream = new SrsStatisticStream();
        stream->vhost = vhost;
        stream->stream = req->stream;
        stream->stream_id = req->stream_id;
        stream->app = req->app;
        stream->url = url;
        rstreams[url] = stream;
        streams[stream->id] = stream;
        return stream;
    }

    stream = rstreams[url];
    return stream;
}
